from __future__ import print_function

import torch
import torch.nn as nn
import torch.nn.functional as F

class SupConLoss(nn.Module):
    """modified supcon loss for segmentation application, the main difference is that the label for different view
    could be different if after spatial transformation"""
    def __init__(self, temperature=0.07, contrast_mode='all',base_temperature=0.07):
        super(SupConLoss, self).__init__()
        self.temperature = temperature
        self.contrast_mode = contrast_mode
        self.base_temperature = base_temperature



    def forward(self, features, labels=None):
        # input features shape: [bsz, v, c, w, h]
        # input labels shape: [bsz, v, w, h]
        device = (torch.device('cuda')
                  if features.is_cuda
                  else torch.device('cpu'))

        if len(features.shape) < 3:
            raise ValueError('`features` needs to be [bsz, n_views, ...],'
                             'at least 3 dimensions are required')
        # input features: [bsz, num_view, c, h ,w], h & w are the image size
        # n_view = 2,,contrast_count=2
        # print("计算机loss contrastive")

        contrast_count = features.shape[1]
        # torch.unbind 返回给定维度所有切片的元组,该维度有多少，就分出多少
        contrast_feature = torch.cat(torch.unbind(features, dim=1), dim=0)  # of size (bsz*v, c, h, w)

        # 将tensor的维度换位。
        kernels = contrast_feature.permute(0, 2, 3, 1) #size (bsz*v,  h, w,c)
        # 每个像素对应长度为contrast_feature.shape[1] 的向量
        kernels = kernels.reshape(-1, contrast_feature.shape[1], 1, 1) #[bsz*v*h*w,c,1,1]
        # kernels = kernels[non_background_idx]

        # 张量和标量做逐元素除法
        # 或者两个可广播的张量之间做逐元素除法
        # F.conv2d(A,B)   A:[batch,inchannel,ih,iw]  B:[outchannel,inchannel/group,kh,kw],group 默认=1
        # print("contrast==",contrast_feature.unique())
        # print("contrast==",contrast_feature.shape)

        logits = torch.div(F.conv2d(contrast_feature, kernels), self.temperature)  # of size (bsz*v, bsz*v*h*w, h, w)
        # print("logits.sum",logits.sum())
        # print("logit.shape",logits.shape)
        logits = logits.permute(1, 0, 2, 3)

        # logits.shape[0] = F.conv2d的输出通道数bsz*v*h*w
        # 第i行的 所有值表示  第i个像素与其他所有像素的分别求内积
        # logits.shape:(bsz*v*h*w,bsz*v*h*w)
        logits = logits.reshape(logits.shape[0], -1)

        if labels is not None:
            # torch.unbind 返回给定维度所有切片的元组
            labels = torch.cat(torch.unbind(labels, dim=1), dim=0)
            #  当调用contiguous()时，会强制拷贝一份tensor，让它的布局和从头创建的一模一样，但是两个tensor完全没有联系。
            labels = labels.contiguous().view(-1, 1)
            # 对两个张量Tensor进行逐元素的比较，若相同位置的两个元素相同，则返回True；若不同，返回False。
            mask = torch.eq(labels, labels.T).float().to(device)

            bg_bool = torch.eq(labels.squeeze().cpu(), torch.zeros(labels.squeeze().shape))
            non_bg_bool = ~ bg_bool
            non_bg_bool = non_bg_bool.int().to(device)
        else:
            # 为了取出锚点

            # 这个函数主要是为了生成对角线全1，其余部分全0的二维数组
            # bsz*v*h*w//contrast_count=2
            mask = torch.eye(logits.shape[0]//contrast_count).float().to(device)
            # mask 对应维度大小扩大 contrast_count倍

            # 当参数只有两个时：（列的重复倍数，行的重复倍数）。1表示不重复
            # 当参数有三个时：（通道数的重复倍数，列的重复倍数，行的重复倍数）
            mask = mask.repeat(contrast_count, contrast_count)
            # print(mask.shape)
        #     mask == tensor([[1., 0., 1., 0.],
                    #         [0., 1., 0., 1.],
                    #         [1., 0., 1., 0.],
                    #         [0., 1., 0., 1.]])

        # mask-out self-contrast cases
        # scatter(input,dim, index, src) 的参数有 3 个
        # dim：沿着哪个维度进行索引
        # index：用来 scatter 的元素索引
        # src：用来 scatter 的源元素，可以是一个标量或一个张量
        # 简单说就是通过一个张量 src  来修改另一个张量，哪个元素需要修改、用 src 中的哪个元素来修改由 dim 和 index 决定
        #  输出的size==input.size

        # ones_lisk 生成形状相同，全1的张量
        # torch.arange(1,3):1,2 不包括3。 torch.range(1,3):1，2，3包括3
        # mask.shape:(bsz*v*h*w,bsz*v*h*w)
        # torch.arange(mask.shape[0]).view(-1, 1) = 0，1，2，3,...,bsz*v*h*w-1
        logits_mask = torch.scatter(
            torch.ones_like(mask),
            1, # 按列，即index的值作为列值索引 一般可以用来对标签进行one-hot 编码.
            torch.arange(mask.shape[0]).view(-1, 1).to(device),
            0 # src的值都是0，用0来填充输出的位置
        )
        # logits_mask的输出应该是对角线全0的张量，类似
        # tensor([[0., 1., 1., 1.],
        #         [1., 0., 1., 1.],
        #         [1., 1., 0., 1.],
        #         [1., 1., 1., 0.]])

        mask = mask * logits_mask

        # compute log_prob
        # logits.shape[0] = F.conv2d的输出通道数bsz*v*h*w
        # 第i行的 所有值表示  第i个像素与其他所有像素的分别求内积
        # logits.shape:(bsz*v*h*w,bsz*v*h*w)

        # 把自己跟自己内积排除
        # print("logits==",logits.unique())
        # print("logits.shape",logits.shape)
        exp_logits = torch.exp(logits) * logits_mask
        # print("exp_logits==",exp_logits.unique())

        # 把第1维串在一起求和(每一行求和)
        # logits 似乎没有经过log的计算？与公式不符？？
        log_prob = logits - torch.log(exp_logits.sum(1, keepdim=True))

        # compute mean of log-likelihood over positive
        mean_log_prob_pos = (mask * log_prob).sum(1) / mask.sum(1)

        # loss
        loss = - mean_log_prob_pos
        # loss = loss.view(anchor_count, batch_size).mean()
        if labels is not None:
            # only consider the contrastive loss for non-background pixel
            loss = (loss * non_bg_bool).sum() / (non_bg_bool.sum())
        else:
            loss = loss.mean()

        return loss


class SupConSegLoss(nn.Module):
    # TODO: only support batch size = 1
    def __init__(self, temperature=0.7):
        super(SupConSegLoss, self).__init__()
        self.temp = temperature
        self.device = (torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu'))

    def forward(self, features, labels=None):
        # input features: [bsz, c, h ,w], h & w are the image size
        shape = features.shape
        img_size = shape[-1]
        if labels is not None:
            f1, f2 = torch.split(features, [1, 1], dim=1)
            features = torch.cat([f1.squeeze(1), f2.squeeze(1)], dim=0)
            l1, l2 = torch.split(labels, [1, 1], dim=1)
            labels = torch.cat([l1.squeeze(1), l2.squeeze(1)], dim=0)
            # features = features.squeeze(dim=1)
            # labels = labels.squeeze(dim=1)
            bsz = features.shape[0]
            loss = []
            for b in range(bsz):
                # print("Iteration index:", idx, "Batch_size:", b)
                for i in range(img_size):
                    # print("before ith iteration, the consumption memory is:", torch.cuda.memory_allocated() / 1024**2)
                    for j in range(img_size):
                        x = features[b:b + 1, :, i:i + 1, j:j + 1]  # [1,c, 1, 1, 1]
                        x_label = labels[b, i, j] + 1  # avoid cases when label=0
                        if x_label == 1:  # ignore background
                            continue
                        cos_dst = F.conv2d(features, x)  # [2b, 1, 512, 512]
                        cos_dst = torch.div(cos_dst.squeeze(dim=1), self.temp)
                        # print("cos_dst:", cos_dst.max(), cos_dst.min())
                        self_contrast_dst = torch.div((x * x).sum(), self.temp)

                        mask = labels + 1
                        mask[mask != x_label] = 0
                        # if mask.sum() < 5:
                        #    print("Not enough same label pixel")
                        #    continue
                        mask = torch.div(mask, x_label)
                        numerator = (mask * cos_dst).sum() - self_contrast_dst
                        denominator = torch.exp(cos_dst).sum() - torch.exp(self_contrast_dst)
                        # print("denominator:", denominator.item())
                        # print("numerator:", numerator.max(), numerator.min())
                        loss_tmp = torch.log(denominator) - numerator / (mask.sum() - 1)
                        if loss_tmp != loss_tmp:
                            print(numerator.item(), denominator.item())

                        loss.append(loss_tmp)
            if len(loss) == 0:
                loss = torch.tensor(0).float().to(self.device)
                return loss
            loss = torch.stack(loss).mean()
            return loss

        else:
            bsz = features.shape[0]
            loss = []
            for b in range(bsz):
                # print("Iteration index:", idx, "Batch_size:", b)
                tmp_feature = features[b]
                for n in range(tmp_feature.shape[0]):
                    for i in range(img_size):
                        # print("before ith iteration, the consumption memory is:", torch.cuda.memory_allocated() / 1024**2)
                        for j in range(img_size):
                            x = tmp_feature[n:n+1, :, i:i + 1, j:j + 1]  # [c, 1, 1, 1]
                            cos_dst = F.conv2d(tmp_feature, x)  # [2b, 1, 512, 512]
                            cos_dst = torch.div(cos_dst.squeeze(dim=1), self.temp)
                            # print("cos_dst:", cos_dst.max(), cos_dst.min())
                            self_contrast_dst = torch.div((x * x).sum(), self.temp)

                            mask = torch.zeros((tmp_feature.shape[0], tmp_feature.shape[2], tmp_feature.shape[3]),
                                               device=self.device)
                            mask[0:tmp_feature.shape[0], i, j] = 1
                            numerator = (mask * cos_dst).sum() - self_contrast_dst
                            denominator = torch.exp(cos_dst).sum() - torch.exp(self_contrast_dst)
                            # print("numerator:", numerator.max(), numerator.min())
                            loss_tmp = torch.log(denominator) - numerator / (mask.sum() - 1)
                            if loss_tmp != loss_tmp:
                                print(numerator.item(), denominator.item())

                            loss.append(loss_tmp)

            loss = torch.stack(loss).mean()
            return loss


class LocalConLoss(nn.Module):
    def __init__(self, temperature=0.7, stride=4):
        super(LocalConLoss, self).__init__()
        self.temp = temperature
        self.device = (torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu'))
        self.supconloss = SupConLoss(temperature=self.temp)
        self.stride = stride

    def forward(self, features, labels=None):

        # input features: [bsz, num_view, c, h ,w], h & w are the image size
        # ::num  表示从0开始间隔num个取值
        features = features[:, :, :, ::self.stride, ::self.stride]  # resample feature maps to reduce memory consumption and running time
        shape = features.shape
        img_size = shape[-1]
        if labels is not None:
            # 间隔取像素
            labels = labels[:, :, ::self.stride, ::self.stride]
            if labels.sum() == 0:
                loss = torch.tensor(0).float().to(self.device)
                return loss

            loss = self.supconloss(features, labels)
            """
            f1, f2 = torch.split(features, [1, 1], dim=1)
            features = torch.cat([f1.squeeze(1), f2.squeeze(1)], dim=0)
            l1, l2 = torch.split(labels, [1, 1], dim=1)
            labels = torch.cat([l1.squeeze(1), l2.squeeze(1)], dim=0)
            bsz = features.shape[0]
            loss = []
            for b in range(bsz):
                # print("Iteration index:", idx, "Batch_size:", b)
                for i in range(img_size):
                    # print("before ith iteration, the consumption memory is:", torch.cuda.memory_allocated() / 1024**2)
                    for j in range(img_size):
                        x = features[b:b + 1, :, i:i + 1, j:j + 1]  # [c, 1, 1, 1]
                        x_label = labels[b, i, j] + 1  # avoid cases when label=0
                        if x_label == 1:  # ignore background
                            continue
                        cos_dst = F.conv2d(features, x)  # [2b, 1, 512, 512]
                        cos_dst = torch.div(cos_dst.squeeze(dim=1), self.temp)
                        self_contrast_dst = torch.div((x * x).sum(), self.temp)

                        mask = labels + 1
                        mask[mask != x_label] = 0
                        mask = torch.div(mask, x_label)
                        numerator = (mask * cos_dst).sum() - self_contrast_dst
                        denominator = torch.exp(cos_dst).sum() - torch.exp(self_contrast_dst)
                        # print("denominator:", denominator.item())
                        # print("numerator:", numerator.max(), numerator.min())
                        loss_tmp = torch.log(denominator) - numerator / (mask.sum() - 1)
                        if loss_tmp != loss_tmp:
                            print(numerator.item(), denominator.item())

                        loss.append(loss_tmp)

            if len(loss) == 0:
                loss = torch.tensor(0).float().to(self.device)
                return loss
            loss = torch.stack(loss).mean()
            """
            return loss
        else:
            bsz = features.shape[0]
            loss = self.supconloss(features)

            """
            loss = []
            for b in range(bsz):
                # print("Iteration index:", idx, "Batch_size:", b)
                tmp_feature = features[b]
                for n in range(tmp_feature.shape[0]):
                    for i in range(img_size):
                        # print("before ith iteration, the consumption memory is:", torch.cuda.memory_allocated() / 1024**2)
                        for j in range(img_size):
                            x = tmp_feature[n:n+1, :, i:i + 1, j:j + 1]  # [c, 1, 1, 1]
                            cos_dst = F.conv2d(tmp_feature, x)  # [2b, 1, 512, 512]
                            cos_dst = torch.div(cos_dst.squeeze(dim=1), self.temp)
                            # print("cos_dst:", cos_dst.max(), cos_dst.min())
                            self_contrast_dst = torch.div((x * x).sum(), self.temp)

                            mask = torch.zeros((tmp_feature.shape[0], tmp_feature.shape[2], tmp_feature.shape[3]),
                                               device=self.device)
                            mask[0:tmp_feature.shape[0], i, j] = 1
                            numerator = (mask * cos_dst).sum() - self_contrast_dst
                            denominator = torch.exp(cos_dst).sum() - torch.exp(self_contrast_dst)
                            # print("numerator:", numerator.max(), numerator.min())
                            loss_tmp = torch.log(denominator) - numerator / (mask.sum() - 1)
                            if loss_tmp != loss_tmp:
                                print(numerator.item(), denominator.item())

                            loss.append(loss_tmp)

            loss = torch.stack(loss).mean()
            """
            return loss


class BlockConLoss(nn.Module):
    def __init__(self, temperature=0.7, block_size=32):
        super(BlockConLoss, self).__init__()
        self.block_size = block_size
        self.device = (torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu'))
        self.supconloss = SupConLoss(temperature=temperature)

    def forward(self, features, labels=None):
        # input features: [bsz, num_view, c, h ,w], h & w are the image size
        shape = features.shape
        img_size = shape[-1]
        div_num = img_size // self.block_size
        if labels is not None:
            loss = []
            for i in range(div_num):
                # print("Iteration index:", idx, "Batch_size:", b)
                for j in range(div_num):
                    # print("before ith iteration, the consumption memory is:", torch.cuda.memory_allocated() / 1024**2)
                    block_features = features[:, :, :, i*self.block_size:(i+1)*self.block_size,
                                  j*self.block_size:(j+1)*self.block_size]
                    block_labels = labels[:,:, i*self.block_size:(i+1)*self.block_size,
                                  j*self.block_size:(j+1)*self.block_size]

                    if block_labels.sum() == 0:
                        continue

                    tmp_loss = self.supconloss(block_features, block_labels)

                    loss.append(tmp_loss)

            if len(loss) == 0:
                loss = torch.tensor(0).float().to(self.device)
                return loss
            loss = torch.stack(loss).mean()
            return loss

        else:
            loss = []
            for i in range(div_num):
                # print("Iteration index:", idx, "Batch_size:", b)
                for j in range(div_num):
                    # print("before ith iteration, the consumption memory is:", torch.cuda.memory_allocated() / 1024**2)
                    block_features = features[:, :, :, i * self.block_size:(i + 1) * self.block_size,
                                     j * self.block_size:(j + 1) * self.block_size]

                    tmp_loss = self.supconloss(block_features)

                    loss.append(tmp_loss)

            loss = torch.stack(loss).mean()
            return loss